import yaml
import paramiko
from prometheus_client import start_http_server, Gauge
import socket
import time
import argparse
import logging
import concurrent.futures

# Define the Prometheus Gauge
TCP_CONNECTIONS = Gauge('tcp_connections', 'TCP connection status', ['src', 'dst', 'port', 'desc'])

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Function to check TCP connection locally
def check_tcp_connection_locally(dst, port):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(1)
    try:
        sock.connect((dst, port))
        logging.info(f"Local connection to {dst}:{port} successful.")
        return True
    except Exception as e:
        logging.error(f"Local connection to {dst}:{port} failed: {e}")
        return False
    finally:
        sock.close()

# Function to check TCP connection via SSH
def check_tcp_connection_via_ssh(src, dst, port, username, key_filename, ssh_port=22):
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        ssh.connect(src, username=username, key_filename=key_filename, port=ssh_port, timeout=2)
        command = f'(echo > /dev/tcp/{dst}/{port}) &>/dev/null && echo "success" || echo "failure"'
        stdin, stdout, stderr = ssh.exec_command(command)
        result = stdout.read().decode().strip()
        if result == "success":
            logging.info(f"SSH connection from {src} to {dst}:{port} via port {ssh_port} successful.")
        else:
            logging.error(f"SSH connection from {src} to {dst}:{port} via port {ssh_port} failed.")
        return result == "success"
    except paramiko.AuthenticationException:
        logging.error(f"Authentication failed when connecting to {src} via SSH port {ssh_port}")
        return False
    except paramiko.SSHException as e:
        logging.error(f"SSHException: {e} when connecting to {src} via SSH port {ssh_port}")
        return False
    except Exception as e:
        logging.error(f"Failed to connect to {src} via SSH on port {ssh_port}: {e}")
        return False
    finally:
        ssh.close()

# Function to load configuration from YAML file
def load_config(config_file):
    try:
        with open(config_file, 'r') as f:
            return yaml.safe_load(f)
    except Exception as e:
        logging.error(f"Failed to load configuration file: {e}")
        raise

# Function to monitor a single server
def monitor_server(server):
    src_addresses = server['src']
    dst_addresses = server['dst']
    dst_ports = server['dstport']
    desc = server['description']

    for src in src_addresses:
        for dst in dst_addresses:
            for port in dst_ports:
                if src == '127.0.0.1':
                    if check_tcp_connection_locally(dst, port):
                        TCP_CONNECTIONS.labels(src=src, dst=dst, port=port, desc=desc).set(1)
                    else:
                        TCP_CONNECTIONS.labels(src=src, dst=dst, port=port, desc=desc).set(0)
                else:
                    ssh_user = server['ssh_user']
                    ssh_key = server['ssh_key']
                    ssh_port = server.get('ssh_port', 22)  # Default SSH port is 22 if not specified
                    if check_tcp_connection_via_ssh(src, dst, port, ssh_user, ssh_key, ssh_port):
                        TCP_CONNECTIONS.labels(src=src, dst=dst, port=port, desc=desc).set(1)
                    else:
                        TCP_CONNECTIONS.labels(src=src, dst=dst, port=port, desc=desc).set(0)

if __name__ == '__main__':
    # Parse command-line arguments
    parser = argparse.ArgumentParser(description='TCP connection checker with Prometheus exporter.')
    parser.add_argument('--port', type=int, default=8000, help='Port to run the Prometheus HTTP server on.')
    parser.add_argument('--config', type=str, default='config.yaml', help='Path to the configuration file.')
    args = parser.parse_args()

    # Start the Prometheus HTTP server
    start_http_server(args.port)
    logging.info(f"Prometheus HTTP server started on port {args.port}.")

    # Load configuration from file
    try:
        config = load_config(args.config)
    except Exception as e:
        logging.critical(f"Failed to load configuration: {e}")
        exit(1)

    # Main monitoring loop
    while True:
        with concurrent.futures.ThreadPoolExecutor() as executor:
            futures = [executor.submit(monitor_server, server) for server in config['servers']]
            concurrent.futures.wait(futures)
        
        time.sleep(10)
